from fastapi import APIRouter, HTTPException
from pydantic import BaseModel
import os
from fastapi.responses import JSONResponse
import random
import string
import pymysql
from config import mysql_config

router_module_real = APIRouter()

@router_module_real.get("/")
def index():
    return {"message": "router_module!"}

# MySQL配置模型
class MySQLConfig(BaseModel):
    host: str
    port: int
    user: str
    password: str
    database: str
    charset: str

# MySQL执行SQL模型
class SQLRequest(BaseModel):
    sql: str

# MySQL使用数据库模型
class UseDatabase(BaseModel):
    database: str

# 当前连接的数据库名
current_database = mysql_config.get("database", "")

# 获取MySQL配置
@router_module_real.get("/mysql/config")
async def get_mysql_config():
    try:
        return {
            "success": True,
            "data": mysql_config
        }
    except Exception as e:
        return {
            "success": False,
            "message": str(e)
        }

# 保存MySQL配置
@router_module_real.post("/mysql/config")
async def save_mysql_config(config: MySQLConfig):
    try:
        global mysql_config
        # 更新全局配置
        mysql_config.update(config.dict())
        
        # 更新config.py文件
        with open("config.py", "r", encoding="utf-8") as f:
            content = f.read()
        
        # 找到mysql_config字典的起始位置
        start_index = content.find("mysql_config = {")
        if start_index == -1:
            # 如果找不到配置，在文件末尾添加
            content += "\n\n# MySQL数据库配置\n"
            content += "mysql_config = {\n"
            content += f'    "host": "{config.host}",\n'
            content += f'    "port": {config.port},\n'
            content += f'    "user": "{config.user}",\n'
            content += f'    "password": "{config.password}",\n'
            content += f'    "database": "{config.database}",\n'
            content += f'    "charset": "{config.charset}"\n'
            content += "}"
        else:
            # 找到配置块的结束位置
            end_index = content.find("}", start_index)
            if end_index == -1:
                raise ValueError("无法解析配置文件")
            
            # 构建新的配置块
            new_config = "mysql_config = {\n"
            new_config += f'    "host": "{config.host}",\n'
            new_config += f'    "port": {config.port},\n'
            new_config += f'    "user": "{config.user}",\n'
            new_config += f'    "password": "{config.password}",\n'
            new_config += f'    "database": "{config.database}",\n'
            new_config += f'    "charset": "{config.charset}"\n'
            new_config += "}"
            
            # 替换配置块
            content = content[:start_index] + new_config + content[end_index+1:]
        
        # 写入配置文件
        with open("config.py", "w", encoding="utf-8") as f:
            f.write(content)
        
        # 更新当前数据库
        global current_database
        current_database = config.database
        
        return {
            "success": True,
            "message": "配置已保存"
        }
    except Exception as e:
        return {
            "success": False,
            "message": str(e)
        }

# 测试数据库连接
@router_module_real.post("/mysql/test-connection")
async def test_connection(config: MySQLConfig):
    try:
        connection = pymysql.connect(
            host=config.host,
            port=config.port,
            user=config.user,
            password=config.password,
            charset=config.charset
        )
        connection.close()
        return {
            "success": True,
            "message": "连接成功"
        }
    except Exception as e:
        return {
            "success": False,
            "message": str(e)
        }

# 获取当前使用的数据库
@router_module_real.get("/mysql/current-db")
async def get_current_db():
    try:
        return {
            "success": True,
            "data": current_database
        }
    except Exception as e:
        return {
            "success": False,
            "message": str(e)
        }

# 获取所有数据库
@router_module_real.get("/mysql/databases")
async def get_databases():
    try:
        connection = pymysql.connect(
            host=mysql_config["host"],
            port=mysql_config["port"],
            user=mysql_config["user"],
            password=mysql_config["password"],
            charset=mysql_config["charset"]
        )
        
        try:
            with connection.cursor() as cursor:
                cursor.execute("SHOW DATABASES")
                databases = [row[0] for row in cursor.fetchall()]
                return {
                    "success": True,
                    "data": databases
                }
        finally:
            connection.close()
    except Exception as e:
        return {
            "success": False,
            "message": str(e)
        }

# 获取当前数据库的所有表
@router_module_real.get("/mysql/tables")
async def get_tables():
    try:
        if not current_database:
            return {
                "success": False,
                "message": "未选择数据库"
            }
        
        connection = pymysql.connect(
            host=mysql_config["host"],
            port=mysql_config["port"],
            user=mysql_config["user"],
            password=mysql_config["password"],
            database=current_database,
            charset=mysql_config["charset"]
        )
        
        try:
            with connection.cursor() as cursor:
                cursor.execute("SHOW TABLES")
                tables = [row[0] for row in cursor.fetchall()]
                return {
                    "success": True,
                    "data": tables
                }
        finally:
            connection.close()
    except Exception as e:
        return {
            "success": False,
            "message": str(e)
        }

# 执行SQL语句
@router_module_real.post("/mysql/execute")
async def execute_sql(request: SQLRequest):
    try:
        # 配置数据库连接参数
        db_config = {
            "host": mysql_config["host"],
            "port": mysql_config["port"],
            "user": mysql_config["user"],
            "password": mysql_config["password"],
            "charset": mysql_config["charset"]
        }
        
        # 如果有当前数据库，则使用它
        if current_database:
            db_config["database"] = current_database
        
        connection = pymysql.connect(**db_config)
        
        try:
            with connection.cursor(pymysql.cursors.DictCursor) as cursor:
                # 执行SQL语句
                affected_rows = cursor.execute(request.sql)
                
                # 判断是否为查询语句
                is_query = request.sql.strip().upper().startswith(("SELECT", "SHOW", "DESCRIBE", "DESC", "EXPLAIN"))
                
                if is_query:
                    # 查询语句，返回结果集
                    result = cursor.fetchall()
                    return {
                        "success": True,
                        "isQuery": True,
                        "data": result
                    }
                else:
                    # 非查询语句，提交事务并返回影响行数
                    connection.commit()
                    return {
                        "success": True,
                        "isQuery": False,
                        "affectedRows": affected_rows
                    }
        finally:
            connection.close()
    except Exception as e:
        return {
            "success": False,
            "message": str(e)
        }

# 使用指定的数据库
@router_module_real.post("/mysql/use-database")
async def use_database(request: UseDatabase):
    try:
        global current_database
        current_database = request.database
        
        # 测试连接
        connection = pymysql.connect(
            host=mysql_config["host"],
            port=mysql_config["port"],
            user=mysql_config["user"],
            password=mysql_config["password"],
            database=current_database,
            charset=mysql_config["charset"]
        )
        connection.close()
        
        return {
            "success": True,
            "message": f"已切换到数据库 {current_database}"
        }
    except Exception as e:
        current_database = ""  # 连接失败，重置当前数据库
        return {
            "success": False,
            "message": str(e)
        }